package edu.gatech.cc.liam.core.rl.models.stochasticgame;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;

import edu.gatech.cc.liam.core.rl.models.Action;


public class ConstructedMultiagentRewardFunction implements RewardFunction<String>{

	private static final long serialVersionUID = -2615080415299813040L;
	
	final protected HashMap<StateJointActionState, Double[]> specificRewardsTemp;
	
	protected HashMap<StateJointActionState, ArrayList<Double>> specificRewardsFinal;
	protected ArrayList<Double> theZeroArray;
	
	int numPlayers;
	Double[] maxRewards;
	Double[] minRewards;
	
	public ConstructedMultiagentRewardFunction(int numPlayers) {
		super();
		this.specificRewardsTemp = new HashMap<StateJointActionState, Double[]>();
		this.numPlayers = numPlayers;
		maxRewards = new Double[numPlayers];
		Arrays.fill(maxRewards, Double.MIN_VALUE);
		minRewards = new Double[numPlayers];
		Arrays.fill(minRewards, Double.MAX_VALUE);
	}

	@SuppressWarnings("unchecked")
	@Override
	public ArrayList<Double> reward(String oldState, ArrayList<Action> action, String newState) {
		StateJointActionState stateJointActionStateKey = new StateJointActionState(oldState, action, newState);
		if(specificRewardsFinal.containsKey(stateJointActionStateKey)) {
			return specificRewardsFinal.get(stateJointActionStateKey);
		}
		return  (ArrayList<Double>)theZeroArray.clone();
	}

	// make all rewards range from exactly 0.0 to 1.0
	public void normalizeTempRewards() {
		double[] max = new double[this.numPlayers];
		Arrays.fill(max, Double.MIN_VALUE);
		double[] min = new double[this.numPlayers];
		Arrays.fill(min, Double.MAX_VALUE);
		for(StateJointActionState aj : specificRewardsTemp.keySet()) {
			Double[] d = specificRewardsTemp.get(aj);
			assert d.length == numPlayers;
			for(int i=0; i<numPlayers; i++) {
				max[i] = Math.max(d[i], max[i]);
				min[i] = Math.min(d[i], min[i]);
			}
		}
		for(StateJointActionState aj : specificRewardsTemp.keySet()) {
			Double[] d = specificRewardsTemp.get(aj);
			for(int i=0; i<numPlayers; i++) {
				d[i] = (d[i] - min[i]) / (max[i]-min[i]);
			}
		}
	}
	
	// add rewards together and fill in wildcards
	public void finalizeRewards(DiscreteActionTransitionFunction<String> aTransFunc) {
		// init the zeroArray
		theZeroArray = new ArrayList<Double>(Arrays.asList(new Double[numPlayers]));
		
		specificRewardsFinal = new HashMap<StateJointActionState, ArrayList<Double>>();
		for(String oldState : aTransFunc.getStates()) {
			for(ArrayList<Action> action : aTransFunc.getAllJointsForState(oldState)) {
				for(String newState : aTransFunc.getSuccessorStateSet(oldState, action)) {
					Double[] reward = new Double[numPlayers];
					Arrays.fill(reward, 0.0);
					checkAndAdd(reward, oldState, action, newState);
					checkAndAdd(reward, oldState, action, null);
					checkAndAdd(reward, oldState, null, newState);
					checkAndAdd(reward, oldState, null, null);
					checkAndAdd(reward, null, action, newState);
					checkAndAdd(reward, null, action, null);
					checkAndAdd(reward, null, null, newState);
					checkAndAdd(reward, null, null, null);
					//if(! isZeros(reward)) {
						StateJointActionState stateJointActionStateKey = new StateJointActionState(oldState, action, newState);
						specificRewardsFinal.put(stateJointActionStateKey, new ArrayList<Double>(Arrays.asList(reward)));
					//}
					for(int i=0; i<this.numPlayers; i++) {
						maxRewards[i] = Math.max(maxRewards[i], reward[i]);
						minRewards[i] = Math.min(minRewards[i], reward[i]);
					}
				}
			}
		}
	}
	
	private void checkAndAdd(Double[] theReward, String oldState, ArrayList<Action> action, String newState) {
		StateJointActionState stateJointActionStateKey = new StateJointActionState(oldState, action, newState);
		if(specificRewardsTemp.containsKey(stateJointActionStateKey)) {
			addTo(theReward, specificRewardsTemp.get(stateJointActionStateKey));
		}
	}
	private boolean isZeros(Double[] a) {
		for(int i=0; i<a.length; i++) {
			if (a[i] != 0.0)
				return false;
		}
		return true;
	}
	private void addTo(Double[] destArray, Double[] addArray) {
		for(int i=0; i<destArray.length; i++) {
			destArray[i] += addArray[i];
		}
	}
	
/**
 * Add a reward to the reward function. Entering a state or jointAction as null will act a wild card.
 * If more than one reward has been entered (because of wild cards) than they will all be added together.
 * If no rewards are given than a reward of zero is given to all players.
 * 
 * @param fromState
 * @param theJoint
 * @param toState
 * @param theReward
 */
	public void addReward(String fromState, ArrayList<Action> theJoint, String toState, Double[] theReward) {
		this.specificRewardsTemp.put(new StateJointActionState(fromState, theJoint, toState), theReward);
	}
	
	public class StateJointActionState {
		String fromState;
		ArrayList<Action> action;
		String toState;
		
		public StateJointActionState(String fromState,
				ArrayList<Action> action, String toState) {
			super();
			this.fromState = fromState;
			this.action = action;
			this.toState = toState;
		}
		
		@Override
		public boolean equals(Object obj) {
			if (obj == null || !getClass().equals(obj.getClass())) {
		        return false;
		     } else {
		    	 StateJointActionState otherObj = (StateJointActionState)obj;
		    	 return nullEquals(this.fromState, otherObj.fromState) &&
		    	 		nullEquals(this.action, otherObj.action) &&
		    	 		nullEquals(this.toState, otherObj.toState);
		     }
		}
		
		private boolean nullEquals(Object o1, Object o2) {
	    	 if (o1 == null)
	    		 return o2 == null;
	    	 return o1.equals(o2);
		}
		
		@Override
		public int hashCode() {
			int hashcode = 0;
			hashcode += fromState == null ? 0 : fromState.hashCode();
			hashcode += action == null ? 0 : action.hashCode();
			hashcode += toState == null ? 0 : toState.hashCode();
			return hashcode;
		}

		@Override
		public String toString() {
			String rStr = "\n";
			rStr += "||" + this.fromState + ", " + this.action + ", " + this.toState + "||";
			return rStr;
		}
		
	}

// doesn't display implicit zero rewards
	@Override
	public String toString() {
		// TODO Auto-generated method stub
		return specificRewardsFinal.toString();
	}
	

	@Override
	public Double[] getMaxRewards() {
		return this.maxRewards;
	}
	@Override
	public Double[] getMinRewards() {
		return this.minRewards;
	}
}
